import numpy as np
import math
from threading import Thread
from OpenGL.GLU import *
from OpenGL.GL import *

class FFTTexture:
    def __init__(self, textureResolution = 256):
        self.textureResolution = textureResolution
        self.fftImage = None
        self.textureID = 0
        self.calculating = False
        self.readyToMakeTexture = False

    def CalculateThreaded(self, sample):
        self.calculating = True
        Thread(target = self.Calculate, args = (sample,)).start()

    def Calculate(self, sample):
        self.calculating = True
        sampleWidth = len(sample)
        step = max(int(sampleWidth / self.textureResolution), 1)

        reds = np.zeros( (self.textureResolution, self.textureResolution) )
        greens = np.zeros( (self.textureResolution, self.textureResolution) )
        blues = np.zeros( (self.textureResolution, self.textureResolution) )

        for i in range(0, sampleWidth, step):
            part = sample[i:i+step]
            transform = np.absolute(np.fft.fft(part, self.textureResolution * 2))
            x = i * float(self.textureResolution) / sampleWidth
            blues[x] = np.square(transform[0:self.textureResolution]) * 16

        greens = np.clip(blues / 16, 0, 255)
        reds = np.clip(blues / 4, 0, 255)
        blues = np.clip(blues, 0, 255)
        blues = np.rot90(blues)
        reds = np.rot90(reds)
        greens = np.rot90(greens)

        img = np.dstack((reds, greens, blues))

        self.fftImage = img
        self.readyToMakeTexture = True
        self.calculating = False

    def GetTexture(self):
        if self.readyToMakeTexture:
            self._MakeTexture()
        return self.textureID

    def _MakeTexture(self):
        glEnable(GL_TEXTURE_2D)
        self.textureID = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, self.textureID)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glPixelStorei(GL_UNPACK_ALIGNMENT,4)
        glTexImage2D(GL_TEXTURE_2D, 0, 3, self.textureResolution, self.textureResolution, 0, GL_RGB, GL_UNSIGNED_BYTE, self.fftImage)
        glDisable(GL_TEXTURE_2D)
        self.readyToMakeTexture = False
        
